#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@File    :   tar.py
@Time    :   2024-05-22 09:45
@Author  :   坐公交也用券
@Version :   1.0
@Contact :   faith01238@hotmail.com
@Homepage : https://liumou.site
@Desc    :   通过下载openssh_arm64_kylin-v10_9.7p-ssl1.1.1f.zip文件并解压替换当前SSH
"""
import logging
import os
import subprocess
import sys
import time
from argparse import ArgumentParser
from datetime import datetime
from subprocess import getoutput, getstatusoutput
from sys import exit

# 创建一个自定义的日志记录器
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# 创建一个自定义的日志处理器，设置其输出格式
formatter = logging.Formatter('%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d | %(funcName)s | %(message)s')
handler = logging.StreamHandler(sys.stdout)
handler.setFormatter(formatter)
logger.addHandler(handler)


class OpenSSHZip:
	def __init__(self, urls, file_name):
		"""
		:param file_name: 文件名
		:param urls: 下载地址
		"""
		self.url = urls
		self.filename = file_name
		self.local_file = False
		# 记录旧版SSH版本
		self.old_version = None
		# 设置安装路径,默认: /usr/local/openssh
		self.ssh_install_dst = "/usr/local/openssh"
		self.ssh_dst_bin = os.path.join(self.ssh_install_dst, "bin")
		self.ssh_dst_sbin = os.path.join(self.ssh_install_dst, "sbin")
		self.ssh_dst_etc = os.path.join(self.ssh_install_dst, "etc")
		# 设置服务名称
		self.service_old_name = 'sshd'
		self.service_new_name = 'sshd'

	def check_os(self):
		"""
		检查操作系统
		:return:
		"""
		res = subprocess.getstatusoutput("cat /etc/os-release| grep ^NAME")
		if res[0] != 0:
			logger.error("无法获取操作系统信息")
			exit(1)
		tmp = res[1].split("=")[1].replace('"', "").lower()
		release_ = tmp
		tmps = tmp.split(" ")
		if len(tmps) > 1:
			release_ = tmps[0]
		# 如果文件名称存在 kylin和arm64则进行系统检查，否则直接返回
		if release_ not in str(self.filename).lower():
			logger.error(f"当前操作系统: {release_} 与 文件名称不匹配,请检查文件名称是否正确")
			exit(1)
		logger.info(f"当前操作系统: {release_} 关键词匹配成功: [ {self.filename} ]")

	# info = "cat /etc/.productinfo  | sed -n 2p"
	# r = getstatusoutput(info)
	# if r[0] == 0:
	# 	if "aarch64-Build20/20210518" in r[1]:
	# 		return True
	# 	else:
	# 		logger.error("当前操作系统不支持")
	# 		print("请使用kylin-v10-rpm-arm64-release-aarch64-Build20/20210518")
	# 		print(f"当前系统信息: {r[1]}")
	# 		print("获取系统信息: cat /etc/.productinfo")
	# 		exit(1)
	# else:
	# 	logger.error("无法获取操作系统信息")
	# 	exit(1)

	def link_ssh_conf(self):
		"""
		链接sshd_config文件
		:return:
		"""
		dst = "/etc/ssh"
		# 获取当前时间,格式如下： 202405201230
		formatted_datetime = datetime.now().strftime('%Y%m%d%H%M%S')
		# 拼接文件名
		conf_new = f"{dst}_{formatted_datetime}"
		# 开始备份
		if os.path.exists(dst):
			c = f"mv {dst} {conf_new}"
			if os.system(c) == 0:
				logger.info(f"备份成功: [ {c} ]")
			else:
				logger.error(f"备份失败: [ {c} ]")
				exit(1)

		os.symlink(self.ssh_dst_etc, dst)

	def delete_init_d(self):
		"""
		删除init.d文件
		:return:
		"""
		fl = ["/etc/init.d/sshd", "/etc/rc.d/init.d/sshd"]
		for file_ in fl:
			if os.path.exists(file_):
				os.remove(file_)

	def service(self, p: str):
		"""
		创建service
		:param p: 可执行文件绝对路径
		:return:
		"""
		file_ = f'/usr/lib/systemd/system/{self.service_new_name}.service'
		# conf = os.path.join(self.ssh_dst_etc, "sshd_config")
		if os.path.exists("/etc/init.d/sshd"):
			os.system('rm -f /etc/init.d/sshd')
		os.system('rm -f /etc/systemd/system/ssh.service')
		os.system('rm -f /etc/systemd/system/sshd.service')
		txt = f"""[Unit]
Description=This is a manually generated SSH service file
After=network.target
StartLimitIntervalSec=5

[Service]
Type=simple
User=root
ExecStart={p}
StandardOutput=journal
KillMode=control-group
Restart=on-failure
RestartPreventExitStatus=255
StandardError=inherit
SyslogIdentifier=sshd

[Install]
WantedBy=multi-user.target
Alias={self.service_new_name}.service\n"""
		try:
			with open(file_, 'w', encoding="utf-8") as fp:
				fp.write(txt)
		except Exception as e:
			logger.error(e)
			return
		os.system(f"systemctl daemon-reload")
		os.system(f"systemctl enable {self.service_new_name}.service")
		os.system(f"systemctl start {self.service_new_name}.service")
		os.system(f"systemctl status {self.service_new_name}.service")

	def get_version(self):
		"""
		获取旧版SSH版本
		:return:
		"""
		self.old_version = getoutput("ssh -V")

	def init_arg(self):
		"""
		初始化参数
		:return:
		"""
		# 获取文件名
		self.filename = self.url.split("/")[-1]
		# 设置bin
		self.ssh_dst_bin = os.path.join(self.ssh_install_dst, "bin")
		# 设置sbin
		self.ssh_dst_sbin = os.path.join(self.ssh_install_dst, "sbin")

	def download(self):
		"""
		下载文件
		:return:
		"""
		logger.info("正在下载文件: {}".format(self.filename))
		# 通过wget下载文件
		c = f"wget {self.url} -O {self.filename}"
		# 执行命令并判断是否下载成功，如果失败则终止程序
		if os.system(c) != 0:
			logger.error("下载失败，请检查网络")
			print(c)
			exit(2)
		logger.info(f"下载成功: [ {self.filename} ]")

	def unzip(self):
		"""
		解压文件
		:return:
		"""
		# 开始判断文件是否存在
		if not os.path.exists(self.filename):
			logger.error(f"文件不存在: [ {self.filename} ]")
			exit(2)
		
		# 根据文件扩展名选择解压命令
		if self.filename.endswith('.tar.gz') or self.filename.endswith('.tgz'):
			# 解压tar.gz文件到/usr/local/目录
			c = f"tar -zxf {self.filename} -C /usr/local/"
		else:
			# 解压zip文件到/usr/local/目录并强制覆盖
			c = f"unzip -o {self.filename} -d /usr/local/"
		
		# 执行解压命令
		if os.system(c) != 0:
			logger.error("解压失败")
			exit(2)
		logger.info("解压成功")

	def copy(self):
		"""
		将解压目录下的openssh/bin目录下的文件复制到/bin目录下
		将解压目录下的openssh/sbin目录下的文件复制到/sbin目录下
		"""
		# 判断文件是否存在
		if not os.path.exists(self.ssh_install_dst):
			logger.error(f"目录不存在: [ {self.ssh_install_dst} ]")
			exit(2)
		# 先获取 /usr/local/openssh/bin 目录文件列表，然后依次查找旧命令所在目录并删除替换
		bin_list = os.listdir(self.ssh_dst_bin)
		for i in bin_list:
			# 通过which查询命令所在目录并删除和替换
			src = os.path.join(self.ssh_dst_bin, i)
			c = f"which {i} | xargs rm -f && cp {src} /bin/"
			if os.system(c) != 0:
				print(c)
				logger.error("复制失败")
				exit(2)
			else:
				logger.info(f"复制成功: [ {src} -> /bin/{i} ]")

		# 先查找sshd命令并删除
		dst_list = []
		while True:
			c = f"which sshd"
			wh = getstatusoutput(c)
			if wh[0] != 0:
				break
			cr = f"rm -f {wh[1]}"
			dst_list.append(wh[1])
			if os.system(cr) != 0:
				print(cr)
				logger.error(f"[ {wh[1]} ] 删除失败")
				exit(2)
		if len(dst_list) > 1:
			for i in dst_list:
				# 拼接源路径
				src = os.path.join(self.ssh_dst_sbin, "sshd")
				c = f"ln -s {src} {i}"
				if os.system(c) != 0:
					print(c)
					logger.error(f"软连接创建失败: [ {i} ]")
					exit(2)
				else:
					logger.info(f"软连接创建成功: [ {i} ]")
		else:
			# 拼接源路径
			src = os.path.join(self.ssh_dst_sbin, "sshd")
			c = f"ln -s {src} /usr/sbin/sshd"
			if os.system(c) != 0:
				print(c)
				logger.error("软连接创建失败")
				exit(2)
			else:
				logger.info("软连接创建成功: [ /usr/sbin/sshd ]")
		logger.info("软连接全部配置完成")

	def restart(self):
		"""
		重启SSH服务
		:return:
		"""
		if os.system("sshd -t") != 0:
			logger.error("SSH配置文件校验失败")
			exit(2)
		# 强制删除 /etc/init.d/sshd
		self.delete_init_d()
		# 重载SSH服务
		c = "systemctl daemon-reload"
		if os.system(c) != 0:
			logger.error(f"服务重载失败: {c}")
			exit(2)
		c = f"systemctl restart {self.service_new_name}"
		if os.system(c) != 0:
			logger.error(f"服务重启失败: [ {self.service_new_name} ]")
			if os.system(c) != 0:
				logger.error(f"服务重启失败: [ {self.service_new_name} ]")
				exit(2)
			else:
				logger.info(f"服务重启成功: [ {self.service_new_name} ]")
			exit(2)
		else:
			logger.info(f"服务重启成功: [ {self.service_new_name} ]")

	def version(self):
		"""
		查看版本信息
		:return:
		"""
		# 获取新版SSH版本
		new_version = getoutput("ssh -V")
		print(f"原版SSH版本: [ {self.old_version} ]")
		print(f"新版SSH版本: [ {new_version} ]")

	def status(self):
		c = f"systemctl status {self.service_new_name}"
		os.system(c)

	def open_root(self):
		"""
		开启root登录
		:return:
		"""
		# 检查是否存在root登录参数选项
		# 拼接源路径
		sshd_config = os.path.join(self.ssh_install_dst, "etc/sshd_config")
		if not os.path.exists(sshd_config):
			logger.error(f"配置文件不存在: [ {sshd_config} ]")
			exit(2)
		# 获取文件内容，判断是否存在注释符号
		text = getoutput(f"cat {sshd_config}")
		# 判断PermitRootLogin是否已配置为yes
		if text.find("#PermitRootLogin") != -1:
			# 如果存在则替换
			c = f"sed -i 's/#PermitRootLogin.*/PermitRootLogin yes/g' {sshd_config}"
			if os.system(c) != 0:
				logger.error(f"开启root登录失败[ {sshd_config} ]")
				exit(2)
			else:
				logger.info(f"开启root登录成功,请在完成升级配置后关闭[ {sshd_config} ]")
		else:
			try:
				with open(sshd_config, 'r+', encoding="utf-8") as fp:
					text = fp.read()
					fp.seek(0)
					fp.truncate()
					fp.write(text.replace("#PermitRootLogin", "PermitRootLogin"))
					fp.close()
			except Exception as e:
				logger.error(f"开启root登录失败[ {e} ]")
				print(sshd_config)
				exit(2)

	# c = f"grep PermitRootLogin {sshd_config}| grep yes"
	# if os.system(c) != 0:
	# 	# 如果不存在则添加
	# 	c = f"sed -i 's/PermitRootLogin.*/PermitRootLogin yes/g'  {sshd_config}"
	# 	if os.system(c) != 0:
	# 		logger.error(f"添加root登录失败[ {sshd_config} ]")
	# 		exit(2)
	# 	else:
	# 		logger.info(f"添加root登录成功,请在完成升级配置后关闭[ {sshd_config} ]")
	# else:
	# 	# 如果存在则替换
	# 	c = f"sed -i 's/PermitRootLogin.*/PermitRootLogin yes/g' {sshd_config}"
	# 	if os.system(c) != 0:
	# 		logger.error(f"开启root登录失败[ {sshd_config} ]")
	# 		exit(2)
	# 	else:
	# 		logger.info(f"开启root登录成功,请在完成升级配置后关闭[ {sshd_config} ]")

	def show(self):
		"""
		显示信息
		:return:
		"""
		print(f"当前版本: [ {self.old_version} ]")
		print(f"新版链接: [ {self.url} ]")
		print(f"安装位置: [ {self.ssh_install_dst} ]")

	def check_unzip(self):
		"""
		检查是否安装unzip
		:return:
		"""
		if os.system("which unzip") != 0:
			logger.error("未安装unzip,请先安装")
			exit(2)

	def start(self, local):
		"""

		:param local: 是否使用本地文件
		:return:
		"""
		self.check_os()
		self.get_version()
		self.init_arg()
		self.show()
		if not local:
			self.download()
		self.unzip()
		self.copy()
		# sshd_shell_text()
		self.service(p="/usr/sbin/sshd")
		self.open_root()
		self.link_ssh_conf()
		self.restart()
		self.status()
		self.version()
		logger.info("升级完成")
		# 拼接路径
		etc = os.path.join(self.ssh_install_dst, "etc")
		print(f"当前服务配置文件位置: {etc}")


def use():
	"""
	使用免责声明
	:return:
	"""
	print("""
	免责声明:
	操作系统环境问题由使用者自行维护,此脚本只是个人编写，不代表任何平台或单位
	1. 本脚本仅供学习交流使用,请勿用于非法用途
	2. 脚本作者不对因使用本脚本而导致的任何问题或损失负责,请先自行阅读源码后再决定是否继续
	3. 可接受脚本bug反馈，但是不提供使用指导,请自行判断,不懂的地方可以通过-h参数获取帮助
	4. 如遇到代码问题可自行修改源码或者反馈，但是作者不会保障及时的修复
	5. 下班时间不会回复、处理任何此类脚本的问题
	""")
	while True:
		c = input("是否继续?(y/n): ")
		if c == "y":
			break
		elif c == "n":
			exit(0)
		else:
			print("输入错误,请重新输入")


class CheckSSl3:
	def __init__(self, p: str):
		"""

		:param p: SSL3安装路径
		"""
		# 是否属于ssl3
		self.ssl3 = False
		self.p = p

	def path(self):
		if not os.path.exists(self.p):
			logger.error(f"路径不存在,请先安装Openssl3.x.x版本并正确配置安装目录: [ {self.p} ]")
			exit(2)

	def version(self):
		"""
		查看openssl版本信息
		:return:
		"""
		c = f"{self.p}/bin/openssl version"
		r = getstatusoutput(c)
		if r[0] != 0:
			logger.error(f"获取版本信息失败: [ {c} ]")
			exit(2)
		else:
			# 获取版本信息
			v = r[1].split(" ")[1]
			# 判断是否属于ssl3
			if "3.0" in v:
				self.ssl3 = True
				logger.info(f"当前版本属于SSL3: [ {v} ]")
			else:
				logger.info(f"当前版本不属于SSL3: [ {v} ]")

	def start(self):
		"""
		开始执行
		:return:
		"""
		logger.info(f"检测到启用了-s3参数安装openssl3版本,正在检测OpenSSL3是否正确安装配置")
		self.path()
		self.version()


def create_user():
	"""
	创建sshd用户
	:return:
	"""
	# 先判断是否已存在
	if subprocess.getstatusoutput("id sshd")[0] != 0:
		if os.system("useradd sshd") != 0:
			logger.error("用户创建失败")
			exit(2)


def service(p: str):
	"""
	创建Service
	:param p:
	:return:
	"""
	if not os.path.exists(p):
		logger.error(f"路径[ {p} ]不存在,请先安装OpenSSH或指向正确的路径")
		exit(2)
	os.system("mkdir -p /usr/lib/systemd/system_bak/")
	os.system("cp -rf /usr/lib/systemd/system/ssh*.service /usr/lib/systemd/system_bak/")
	os.system('rm -f /usr/lib/systemd/system/ssh*.service')
	if os.path.exists("/etc/init.d/sshd"):
		os.system('rm -f /etc/init.d/sshd')
	f = '/usr/lib/systemd/system/sshd.service'
	txt = f"""[Unit]
Description=This is a manually generated SSH service file
After=network.target
StartLimitIntervalSec=5

[Service]
Type=simple
User=root
ExecStart={p}
StandardOutput=journal
KillMode=control-group
Restart=on-failure
RestartPreventExitStatus=255
StandardError=inherit
SyslogIdentifier=sshd

[Install]
WantedBy=multi-user.target
Alias=sshd.service"""
	try:
		with open(f, 'w', encoding="utf-8") as fp:
			fp.write(txt)
	except Exception as e:
		logger.error(e)
		exit(2)

	os.system(f"systemctl daemon-reload")
	os.system(f"systemctl enable sshd.service")
	if os.system(f"systemctl restart sshd.service") != 0:
		logger.error("启动失败,正在恢复")
		c = "cp -rf /usr/lib/systemd/system_bak/* /usr/lib/systemd/system/"
		os.system(c)
		os.system(f"systemctl daemon-reload")
		if os.system(f"systemctl restart sshd.service") != 0:
			logger.error("启动失败,请手动启动")
			exit(2)

	os.system(f"systemctl status sshd.service")


def create_service():
	"""

	:return:
	"""
	if not os.path.isfile("/usr/local/openssh/sbin/sshd"):
		return
	logger.info("将仅生成Service文件,如需重新编译请先删除")
	service(p="/usr/sbin/sshd")
	exit(0)


if __name__ == '__main__':
	start_time = time.time()
	local_file = False
	default_url = "https://gitee.com/liumou_site/openssh/releases/download/10.0p1/OpenSSH-10.0p1-KylinV10_Tercel-arm64-OpenSSL3.5.0.tar.gz"
	arg = ArgumentParser(description='当前脚本版本: 1.0_20250102', prog="OpenSSHZip")
	h = f"设置下载地址(默认使用该参数),默认: {default_url}"
	arg.add_argument("-u", "--url", help=h, default=default_url, type=str)
	h = "指定本地文件,不从网页下载,例如：openssh_arm64.zip"
	arg.add_argument("-f", "--file", help=h, default=None, type=str)
	# 使用布尔值判断是否需要创建service
	arg.add_argument("-q", "--quit", help=f"创建service之后直接退出", action="store_true")
	# 解析参数
	args = arg.parse_args()
	# 获取URL值
	url = args.url
	# 获取是否使用本地文件
	file = args.file
	quit_arg = arg.parse_args().quit
	if quit_arg:
		create_service()
		exit(0)
	if file is not None:
		if not os.path.exists(file):
			logger.error(f"本地文件不存在: [ {file} ]")
			exit(2)
		else:
			local_file = True
			logger.info(f"使用本地文件: [ {file} ]")
	else:
		logger.info(f"使用默认下载地址: [ {url} ]")
		file = str(url).split("/")[-1]
	os.system("mkdir -p /var/empty")

	use()
	create_user()
	zo = OpenSSHZip(urls=url, file_name=file)
	zo.start(local=local_file)
	end_time = time.time()
	logger.info(f"升级完成,耗时: {end_time - start_time} 秒")
